package com.cloudeggtech.granite.sdk;

import com.cloudeggtech.chalk.core.stream.IAuthenticationCallback;
import com.cloudeggtech.chalk.core.stream.IAuthenticationToken;
import com.cloudeggtech.chalk.core.stream.INegotiationContext;
import com.cloudeggtech.chalk.core.stream.NegotiationException;
import com.cloudeggtech.chalk.core.stream.StandardStreamer;
import com.cloudeggtech.chalk.core.stream.UsernamePasswordToken;
import com.cloudeggtech.chalk.core.stream.negotiants.InitialStreamNegotiant;
import com.cloudeggtech.chalk.core.stream.negotiants.sasl.SaslError;
import com.cloudeggtech.chalk.network.ConnectionException;
import com.cloudeggtech.chalk.utils.Base64;
import com.cloudeggtech.basalt.protocol.oxm.IOxmFactory;
import com.cloudeggtech.basalt.protocol.oxm.OxmService;
import com.cloudeggtech.basalt.protocol.oxm.annotation.AnnotatedParserFactory;
import com.cloudeggtech.basalt.protocol.oxm.parsers.SimpleObjectParserFactory;
import com.cloudeggtech.basalt.protocol.oxm.parsers.core.stream.BindParser;
import com.cloudeggtech.basalt.protocol.oxm.parsers.core.stream.sasl.FailureParserFactory;
import com.cloudeggtech.basalt.protocol.oxm.translators.SimpleObjectTranslatorFactory;
import com.cloudeggtech.basalt.protocol.oxm.translators.core.stream.sasl.AuthTranslatorFactory;
import com.cloudeggtech.basalt.protocol.core.IError;
import com.cloudeggtech.basalt.protocol.core.ProtocolChain;
import com.cloudeggtech.basalt.protocol.core.stream.Bind;
import com.cloudeggtech.basalt.protocol.core.stream.Feature;
import com.cloudeggtech.basalt.protocol.core.stream.Features;
import com.cloudeggtech.basalt.protocol.core.stream.Session;
import com.cloudeggtech.basalt.protocol.core.stream.Stream;
import com.cloudeggtech.basalt.protocol.core.stream.sasl.Abort;
import com.cloudeggtech.basalt.protocol.core.stream.sasl.Auth;
import com.cloudeggtech.basalt.protocol.core.stream.sasl.Challenge;
import com.cloudeggtech.basalt.protocol.core.stream.sasl.Failure;
import com.cloudeggtech.basalt.protocol.core.stream.sasl.Mechanisms;
import com.cloudeggtech.basalt.protocol.core.stream.sasl.Response;
import com.cloudeggtech.basalt.protocol.core.stream.sasl.Success;

import java.io.IOException;
import java.util.Hashtable;
import java.util.List;

import javax.security.auth.callback.Callback;
import javax.security.auth.callback.CallbackHandler;
import javax.security.auth.callback.NameCallback;
import javax.security.auth.callback.PasswordCallback;
import javax.security.auth.callback.UnsupportedCallbackException;
import javax.security.sasl.RealmCallback;
import javax.security.sasl.Sasl;
import javax.security.sasl.SaslClient;
import javax.security.sasl.SaslClientFactory;
import javax.security.sasl.SaslException;

import cryptix.sasl.ClientFactory;
import cryptix.sasl.SaslParams;


public class AndroidSaslNegotiant extends InitialStreamNegotiant {

    private static final int MAX_FAILURE_COUNT_EXCCEED_READ_RESPONSE_TIMEOUT = 200;

    private static final int DEFAULT_SASL_PROCESS_TIMEOUT = 5000 * 2;
    public static final String DIGEST_MD5_MECHANISM = "DIGEST-MD5";
    public static final String CRAM_MD5_MECHANISM = "CRAM-MD5";


    private static IOxmFactory processingFactory = OxmService.createStreamOxmFactory();

    private UsernamePasswordToken authToken;
    private int failureCount;
    private IAuthenticationCallback authCallback;
    private volatile boolean waitFailureAction;
    private boolean abortSasl;
    public static String mechanism = CRAM_MD5_MECHANISM;
    private static SaslClientFactory clientFactory = new ClientFactory();

    static {
        InitialStreamNegotiant.oxmFactory.register(ProtocolChain.first(Features.PROTOCOL).next(Bind.PROTOCOL),
                new AnnotatedParserFactory<Bind>(BindParser.class)
        );
        InitialStreamNegotiant.oxmFactory.register(ProtocolChain.first(Features.PROTOCOL).next(Session.PROTOCOL),
                new SimpleObjectParserFactory<Session>(
                        Session.PROTOCOL,
                        Session.class)
        );

        processingFactory.register(ProtocolChain.first(Challenge.PROTOCOL),
                new SimpleObjectParserFactory<Challenge>(
                        Challenge.PROTOCOL,
                        Challenge.class)
        );
        processingFactory.register(ProtocolChain.first(Failure.PROTOCOL), new FailureParserFactory());
        processingFactory.register(ProtocolChain.first(Success.PROTOCOL),
                new SimpleObjectParserFactory<Success>(
                        Success.PROTOCOL,
                        Success.class)
        );

        processingFactory.register(Auth.class, new AuthTranslatorFactory());
        processingFactory.register(Response.class, new SimpleObjectTranslatorFactory<Response>(
                Response.class,
                Response.PROTOCOL)
        );
        processingFactory.register(Abort.class, new SimpleObjectTranslatorFactory<Abort>(
                Abort.class,
                Abort.PROTOCOL));
    }


    public static void setSaslClientFactory(SaslClientFactory clientFactory) {
        AndroidSaslNegotiant.clientFactory = clientFactory;
    }

    public AndroidSaslNegotiant(String hostName, String lang, IAuthenticationToken authToken) {
        super(hostName, lang);

        checkAuthToken(authToken);

        this.authToken = (UsernamePasswordToken) authToken;
        this.failureCount = 0;

        abortSasl = false;
    }

    private void checkAuthToken(IAuthenticationToken authToken) {
        if (authToken == null) {
            throw new IllegalArgumentException("null authentication token");
        }

        if (!(authToken instanceof UsernamePasswordToken)) {
            throw new IllegalArgumentException("only UsernamePasswordToken supported");
        }
    }

    @Override
    protected void doNegotiate(INegotiationContext context) throws ConnectionException, NegotiationException {
        @SuppressWarnings("unchecked")
        List<Feature> features = (List<Feature>) context.getAttribute(StandardStreamer.NEGOTIATION_KEY_FEATURES);
        Mechanisms mechanisms = findMechanisms(features);

        if (mechanisms != null) {
            if (!mechanismFound(mechanisms)) {
                new NegotiationException(this, SaslError.MECHANISM_NOT_SUPPORTED);
            }

            negotiateSasl(context);

            if (!abortSasl) {
                super.doNegotiate(context);
            }
        }
    }

    protected void negotiateSasl(INegotiationContext context) throws ConnectionException, NegotiationException {
        if (abortSasl || authToken == null)
            return;

        SaslClient saslClient = this.createSaslClient();

        context.write(processingFactory.translate(new Auth(mechanism)));

        while (!saslClient.isComplete()) {
            Object response = processingFactory.parse(readResponse(DEFAULT_SASL_PROCESS_TIMEOUT));
            if (response instanceof Challenge) {
                String saslResponse = getSaslResponse(saslClient, (Challenge) response);
                context.write(processingFactory.translate(new Response(saslResponse)));
            } else if (response instanceof Failure) {
                processFailure(context, (Failure) response);
                break;
            } else if (response instanceof IError) {
                processError((IError) response, context, processingFactory);
            } else {
                //success
                return;
            }
        }
        if (saslClient.isComplete()) {
            Object response = processingFactory.parse(readResponse(DEFAULT_SASL_PROCESS_TIMEOUT));
            if (response instanceof Success) {
                return; // success
            } else if (response instanceof Failure) {
                processFailure(context, (Failure) response);
            } else {
                processError((IError) response, context, processingFactory);
                return;
            }
        }

        try {
            synchronized (this) {
                if (waitFailureAction) {
                    wait();
                }
            }
        } catch (InterruptedException e) {
            throw new NegotiationException(this, e);
        }

        if (abortSasl) {
            context.write(processingFactory.translate(new Abort()));
            Failure failure = (Failure) processingFactory.parse(readResponse(DEFAULT_SASL_PROCESS_TIMEOUT));

            if (failure.getErrorCondition() == Failure.ErrorCondition.ABORTED) {
                context.write(processingFactory.translate(new Stream(true)));
                Stream closeStream = (Stream) processingFactory.parse(readResponse(DEFAULT_SASL_PROCESS_TIMEOUT));

                if (Boolean.TRUE.equals(closeStream.getClose())) {
                    context.close();
                }

                return;
            }
        } else {
            negotiateSasl(context);
        }
    }


    protected SaslClient createSaslClient() throws NegotiationException {
        Hashtable<String, String> props = new Hashtable<String, String>();
        props.put(Sasl.QOP, "auth");
        props.put(SaslParams.USERNAME, this.authToken.getUsername());
        props.put(SaslParams.PASSWORD, new String(this.authToken.getPassword()));
        Sasl.setSaslClientFactory(AndroidSaslNegotiant.clientFactory);
        try {
            SaslClient saslClient = Sasl.createSaslClient(new String[]{mechanism}, this.authToken.getUsername(),
                    com.cloudeggtech.basalt.protocol.Constants.PROTOCOL_NAME, hostName, props, new DefaultCallbackHandler());
            return saslClient;
        } catch (SaslException e) {
            throw new NegotiationException(this, e);
        }
    }

    private void processFailure(INegotiationContext context, Failure response) throws NegotiationException, ConnectionException {
        waitFailureAction = true;
        authToken = null;
        failureCount++;

        Failure failure = response;
        try {
            String message = readResponse(MAX_FAILURE_COUNT_EXCCEED_READ_RESPONSE_TIMEOUT);
            Stream closeStream = (Stream) processingFactory.parse(message);
            if (closeStream.getClose()) {
                context.close();
            }

            authCallback.failed(new SaslAuthenticationFailure(this, failure.getErrorCondition(), false, failureCount));

            throw new NegotiationException(this, SaslError.MAX_FAILURE_COUNT_EXCCEED);
        } catch (ConnectionException e) {
            if (e.getType() == ConnectionException.Type.READ_RESPONSE_TIMEOUT) {
                authCallback.failed(new SaslAuthenticationFailure(this, failure.getErrorCondition(), true, failureCount));
                return;
            }

            throw e;
        }
    }

    private String getSaslResponse(SaslClient saslClient, Challenge challenge) throws NegotiationException {
        byte[] data = Base64.decode(challenge.getText());

        try {
            data = saslClient.evaluateChallenge(data);
        } catch (SaslException e) {
            throw new NegotiationException(this, e);
        }

        if (data == null)
            return null;

        return Base64.encodeToString(data, false);
    }

    private class DefaultCallbackHandler implements CallbackHandler {

        @Override
        public void handle(Callback[] callbacks) throws IOException, UnsupportedCallbackException {
            for (Callback callback : callbacks) {
                if (callback instanceof NameCallback) {
                    NameCallback nameCallback = (NameCallback) callback;
                    nameCallback.setName(authToken.getUsername());
                } else if (callback instanceof PasswordCallback) {
                    PasswordCallback passwordCallback = (PasswordCallback) callback;
                    passwordCallback.setPassword(authToken.getPassword());
                } else if (callback instanceof RealmCallback) {
                    RealmCallback realmCallback = (RealmCallback) callback;
                    realmCallback.setText(hostName);
                }
            }
        }

    }

    private boolean mechanismFound(Mechanisms mechanisms) {
        for (String m : mechanisms.getMechanisms()) {
            if (mechanism.equals(m)) {
                return true;
            }
        }

        return false;
    }

    private Mechanisms findMechanisms(List<Feature> features) {
        for (Feature feature : features) {
            if (feature instanceof Mechanisms)
                return (Mechanisms) feature;
        }

        return null;
    }

    public void setAuthenticationCallback(IAuthenticationCallback authCallback) {
        this.authCallback = authCallback;
    }

    public IAuthenticationCallback getAuthenticationCallback() {
        return authCallback;
    }

    public void retry(IAuthenticationToken authToken) {
        checkAuthToken(authToken);

        this.authToken = (UsernamePasswordToken) authToken;
        synchronized (this) {
            waitFailureAction = false;
            notify();
        }
    }

    public void abort() {
        this.authToken = null;

        synchronized (this) {
            waitFailureAction = false;
            abortSasl = true;
            notify();
        }
    }
}
